訓練程式碼 part.2

import torch
from tqdm.auto import tqdm

progress_bar = tqdm(range(num_training_steps))

# directory to save the models
output_dir = "MLP_TrainedModels"

for epoch in range(num_train_epochs):
    # Training
    model.train()
    for batch in train_dataloader:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

    # Evaluation
    model.eval()
    losses = []
    for step, batch in enumerate(eval_dataloader):
        with torch.no_grad():
            outputs = model(**batch)
        loss = outputs.loss
        losses.append(accelerator.gather(loss.repeat(batch_size)))

    losses = torch.cat(losses)
    losses = losses[: len(eval_dataset)]

    # perplexity metric used for mask language model training
    try:
        perplexity = math.exp(torch.mean(losses))
    except OverflowError:
        perplexity = float("inf")
    print(f">>> Epoch {epoch}: Perplexity: {perplexity}")

    # Save model
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
    if accelerator.is_main_process:
        tokenizer.save_pretrained(output_dir)

首先使用python內建的tqdm設定一個進度條來監控訓練進度，然後為輸出訓練模型設定一個目錄。接下來定義一個for循環來循環遍歷epoch數，對於每個epoch，我們開始資料集的訓練，循環遍歷訓練資料載入器，計算模型的輸出，計算輸出的損失，使用變壓器導入加速器包對模型進行反向傳播，使用優化器優化模型以最小化損失。並且應用學習率調度程序，使用優化器將累積梯度設為零，並更新了進度條。 我們這樣做，直到完成一個時期的整個資料集的訓練。然後在測試資料集上評估了一個時期的訓練模型，計算了測試資料集上的損失，類似於訓練期間的計算方式，計算模型的交叉熵損失，然後計算損失的指數以獲得模型的 Perplexity。最後使用加速器保存預訓練模型，並使用分詞器保存有關模型的重要文件，例如分詞器和詞彙資訊。訓練後的模型和設定檔保存在輸出目錄資料夾 MLP_TrainedModels 中。